from treeKernel import ftree
from classDef import ErrorCode, affmap, rollback_on_exception,format_error
from typing import Dict,Union,Optional
from visitPCfs import visitPCfs
from pydantic import BaseModel, Field
import os
from termcolor import colored
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

#C
class createDocModel(BaseModel):
    space_int: int = Field(..., description="目标空间ID。对于所有空间/文档ID,请至少先调用searchModel看清文件系统，每个doc或space的id是什么再填这个")
    newDocName: str = Field(..., description="新文档名称。加上后缀名(一般为.md)。注意，目前文件系统要求遵循严格的：根空间下必须全是空间(文件夹),而没有文件。每个空间下全是文件，没有子空间")
class createSpaceModel(BaseModel):
    newSpaceName: str = Field(..., description="新空间名称(等同于目录).注意，目前文件系统要求遵循严格的：根空间下必须全是空间(文件夹),而没有文件。每个空间下全是文件，没有子空间")
class moveDocModel(BaseModel):
    doc_int: int = Field(..., description="要移动的文档ID")
    dest_space_int: int = Field(..., description="目标空间ID")
class copyDocModel(BaseModel):
    doc_int: int = Field(..., description="要复制的文档ID")
    dest_space_int: int = Field(..., description="目标空间ID")
#R
class readDocModel(BaseModel):
    doc_int: int = Field(..., description="要读取的文档ID.注意，返回的文档在'|'及之前是行序号，非原文内容")
class searchModel(BaseModel):
    query: str = Field(..., description="要搜索的文本。若无具体搜索内容，仅查看文件，请留空并设置ensure_full_return_type参数。注意，返回搜素到的内容是不连续的文档块，不可以从返回结果的行数推断某个文字处在原文的哪一行中。")
    ensure_type: str = Field(..., description=(
        """本搜索模型基于广义的DOM结构，将文件系统树与Markdown语法树整合为一棵统一的树结构，其中每个节点均具有明确的类型。
        可指定的节点类型包括：space, doc, h1, h2, h3, h4, h5, h6, li, p, blockquote, pre, code, hr, img, a。
        若不指定此参数，返回的节点类型将不固定，仅包括与搜索语义或文字匹配的节点。
        指定此参数后，搜索结果将额外包含所有符合指定类型的节点及其所有祖先节点，使Agent能够全面掌握到指定类型层级的节点信息。
        例如，指定'space'将返回所有space类型的节点及其祖先节点，'doc'则返回所有doc类型的节点及其祖先节点。
        如无特定查看需求，可留空此参数。"""
    ))
    method: str = Field(..., description="搜索方法，可选值为 'regex'（正则表达式）, 'literal'（字面量匹配）, 'semantic'（语义匹配）。")
    subrootid: int = Field(..., description="从某个子节点作为根，开始向下搜索，填它的id。一般应该从真正的根开始搜素(id为1)")
#U
class edit(BaseModel):
    start_line:int = Field(..., description="要修改的起始行数.最小行数为1")
    end_line:int = Field(..., description="要修改的结束行数")
    content:str = Field(..., description="要修改的内容。注意，返回的文档在'|'及之前是行序号，非原文内容，写content时不要包含进去！！")
class editDocModel(BaseModel):
    doc_int: int=Field(..., description="要修改的文件的id.务必注意：在调用editDocModel前，必须必须先调用readDocModel看清全文文本和对应的行数。")
    edits: list[edit]=Field(..., description="要修改的内容的edit对象列表")
class renameDocModel(BaseModel):
    doc_int: int = Field(..., description="要重命名的文档ID")
    newDocName: str = Field(..., description="新文档名称")
class renameSpaceModel(BaseModel):
    space_int: int = Field(..., description="要重命名的空间ID")
    newSpaceName: str = Field(..., description="新空间名称")
#D
class deleteDocModel(BaseModel):
    doc_int: int = Field(..., description="要删除的文档ID")
class deleteSpaceModel(BaseModel):
    space_int: int = Field(..., description="要删除的空间ID")

AGENT_TOOLS=[createDocModel,createSpaceModel,moveDocModel,copyDocModel,readDocModel,searchModel,editDocModel,renameDocModel,renameSpaceModel,deleteDocModel,deleteSpaceModel]


class crudAgent:
    def __init__(self, tree: ftree): 
        self.tree = tree
        self.session=self.tree.session
        self.err = ErrorCode(self.session)
        self.affine = visitPCfs()
    def router(self,user:str,name:str,args)->Dict[str,Union[str,int]]:
        # print(colored(f"crudAgent router: {name} {args} {self.tree.pendingPermissions}","green"))
        #C
        try:
            if name == "createDocModel":
                return self.createDoc(user,int(args["space_int"]),args["newDocName"])
            elif name == "createSpaceModel":
                return self.createSpace(user,args["newSpaceName"])
            elif name == "moveDocModel":
                return self.moveDoc(user,int(args["doc_int"]),int(args["dest_space_int"]))
            elif name == "copyDocModel":
                return self.copyDoc(user,int(args["doc_int"]),int(args["dest_space_int"]))
            #R
            elif name == "readDocModel":
                return self.readDoc(user,int(args["doc_int"]),True,"agent")
            elif name == "searchModel":
                return self.search(user,args["query"],args["ensure_type"],args["method"],args["subrootid"])
            #U
            elif name == "editDocModel":
                return self.editDoc(user,int(args["doc_int"]),args["edits"])
            elif name == "renameDocModel":
                return self.renameDoc(user,int(args["doc_int"]),args["newDocName"])
            elif name == "renameSpaceModel":
                return self.renameSpace(user,args["space_int"],args["newSpaceName"])
            #D
            elif name == "deleteDocModel":
                return self.deleteDoc(user,int(args["doc_int"]))
            elif name == "deleteSpaceModel":
                return self.deleteSpace(user,int(args["space_int"]))
            else:
                return self.err.UNKNOWN_TOOL_CALL(name)
        except Exception as e:
            return self.err.INTERNAL_EXCEPTION(e)


    @rollback_on_exception
    def createDoc(self, user: str, space_int: int, newDocName: str) -> Dict[str, Union[str,int]]:
        if space_int==1:
            return self.err.INVALID_PARAM(f"createDoc under root space is not allowed")
        i_docId_Dict = self.tree.addNode(user, space_int, newDocName, "doc", by="agent")
        i_docId=int(i_docId_Dict.get("success", -1))
        if i_docId < 0:
            return i_docId_Dict
        SpaceMap = self.session.query(affmap).where(affmap.intId == space_int).first()# type: ignore
        if SpaceMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"SpaceMap from createDoc,intId={space_int} is None",space_int)
        docMap = affmap(intId=int(i_docId_Dict.get("success", -1)),
                a_spaceId=SpaceMap.a_spaceId,
                a_docId="CreateDoc_TobeReplaced")
        self.session.add(docMap)
        self.session.flush()
        #Assumption: agent自己加的节点，自己得能看得到?
        a_docId_Dict = self.affine.createDoc(user,SpaceMap.a_spaceId, newDocName)
        docMap.a_docId=a_docId_Dict.get("affDocId", "")
        a_docId_Dict["nodes"]=self.tree.pNodeList([SpaceMap.intId,docMap.intId],style="ja")
        return self.err.IF_SUCCESS_COMMIT_A2I(a_docId_Dict, i_docId_Dict)
    @rollback_on_exception
    def createSpace(self, user: str, newSpaceName: str) -> Dict[str, Union[str,int]]:
        i_spaceId_Dict = self.tree.addNode(user, self.tree.root(), newSpaceName, "space", by="agent")
        if int(i_spaceId_Dict.get("success", -1)) < 0:
            return i_spaceId_Dict
        spaceMap = affmap(intId=int(i_spaceId_Dict.get("success", -1)),
                    a_spaceId="CreateSpace_TobeReplaced",
                    a_docId="")
        self.session.add(spaceMap)
        self.session.flush()
        a_spaceId_Dict = self.affine.createSpace(user,newSpaceName)
        spaceMap.a_spaceId=a_spaceId_Dict.get("affSpaceId", "")
        # (f"spaceMap {spaceMap.intId} {spaceMap.a_spaceId} {spaceMap.a_docId}")
        return self.err.IF_SUCCESS_COMMIT_A2I(a_spaceId_Dict, i_spaceId_Dict)
    def printMap(self):
        for i,spaceMap in enumerate(self.session.query(affmap).all()):
            print(f"{i}:intId={spaceMap.intId},a_spaceId={spaceMap.a_spaceId},a_docId={spaceMap.a_docId}")
    @rollback_on_exception
    def renameDoc(self, user: str, doc_int: int, newDocName: str) -> Dict[str, Union[str,int]]:
        i_docId_Dict = self.tree.renameNode(user, doc_int, newDocName, by="agent")
        if int(i_docId_Dict.get("success", -1)) < 0:
            return i_docId_Dict
        DocMap = self.session.query(affmap).where(affmap.intId == doc_int).first()# type: ignore
        if DocMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"DocMap from renameDoc,intId={doc_int} is None",doc_int)
        a_Dict = self.affine.renameDoc(user,DocMap.a_spaceId, DocMap.a_docId, newDocName)
        DocMap.a_docId=newDocName#to do:if pluged into real affine ,then this should be removed!!! REMOVE_TAG.
        self.session.add(DocMap)#REMOVE_TAG BECAUSE actually , affineDocId will not change by rename , 
        #but if affineDocId and affineSpaceId are used for local filesystem resource location, then DocMap should be refreshed
        a_Dict["nodes"]=self.tree.pNodeList([DocMap.intId],style="ja")
        return self.err.IF_SUCCESS_COMMIT_A2I(a_Dict, i_docId_Dict)
    @rollback_on_exception
    def renameSpace(self, user: str, space_int: int, newSpaceName: str) -> Dict[str, Union[str,int]]:
        i_spaceId_Dict = self.tree.renameNode(user, space_int, newSpaceName, by="agent")
        if int(i_spaceId_Dict.get("success", -1)) < 0:
            return i_spaceId_Dict
        SpaceMap = self.session.query(affmap).where(affmap.intId == space_int).first()# type: ignore
        if SpaceMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"SpaceMap from renameSpace,intId={space_int} is None",space_int)
        a_Dict = self.affine.renameSpace(user,SpaceMap.a_spaceId, newSpaceName)
        #REMOVE_TAG to do:if pluged into real affine ,then this should be removed!!!
        SpaceDocsMap=self.session.query(affmap).where(affmap.a_spaceId==SpaceMap.a_spaceId).all()#type: ignore 
        for docMap in SpaceDocsMap:
            docMap.a_spaceId=a_Dict.get("affSpaceId", "")
        self.session.add_all(SpaceDocsMap)#REMOVE_TAG
        a_Dict["nodes"]=self.tree.pNodeList([SpaceDocMap.intId for SpaceDocMap in SpaceDocsMap],style="ja")#type: ignore
        return self.err.IF_SUCCESS_COMMIT_A2I(a_Dict, i_spaceId_Dict)

    @rollback_on_exception
    def deleteDoc(self, user: str, doc_int: int) -> Dict[str, Union[str,int]]:
        i_docId_Dict = self.tree.deleteNode(user, doc_int, by="agent")
        if int(i_docId_Dict.get("success", -1)) < 0:
            return i_docId_Dict
        DocMap = self.session.query(affmap).where(affmap.intId == doc_int).first()# type: ignore
        if DocMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"DocMap from deleteDoc,intId={doc_int} is None", doc_int)
        self.session.delete(DocMap)
        a_Dict = self.affine.deleteDoc(user,DocMap.a_spaceId, DocMap.a_docId)
        return self.err.IF_SUCCESS_COMMIT_A2I(a_Dict, i_docId_Dict)
    @rollback_on_exception
    def deleteSpace(self, user: str, space_int: int) -> Dict[str, Union[str,int]]:
        if space_int==1:
            return self.err.INVALID_PARAM(f"delete root space is not allowed")
        i_spaceId_Dict = self.tree.deleteNode(user, space_int, by="agent")
        if int(i_spaceId_Dict.get("success", -1)) < 0:
            return i_spaceId_Dict
        SpaceMap = self.session.query(affmap).where(affmap.intId == space_int).first()# type: ignore
        if SpaceMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"SpaceMap from deleteSpace,intId={space_int} is None",space_int)
        self.session.delete(SpaceMap)
        a_Dict = self.affine.deleteSpace(user,SpaceMap.a_spaceId)
        return self.err.IF_SUCCESS_COMMIT_A2I(a_Dict, i_spaceId_Dict)
    @rollback_on_exception
    def editDoc(self, user: str, doc_int: int, edits: list[Dict]) -> Dict[str, Union[str,int]]:
        # 获取当前文档内容
        #the readDoc called here should be "by=system",
        read_result = self.readDoc(user, doc_int,True,"agent")
        if int(read_result.get("success", -1)) < 0:
            return read_result
        content = str(read_result.get("content", ""))
        lines = content.split('\n')
        # 如果文档行号是从1开始，那么使用下面这种方式进行行号转换
        # （start_line - 1）→在Python列表中对应的实际索引
        #
        # 按倒序处理编辑操作是为了在执行多次替换时
        # 不影响后续edit中尚未处理的行号索引
        for edit in sorted(edits, key=lambda x: x['start_line'], reverse=True):
            start_index = edit['start_line'] - 1
            end_index = edit['end_line'] - 1

            # 校验行号范围
            if start_index < 0 or end_index >= len(lines):
                return self.err.INVALID_PARAM(
                    f"invalid line number range: {edit['start_line']}-{edit['end_line']} (total lines: {len(lines)})"
                )

            # 替换指定行内容
            new_lines = edit['content'].split('\n')
            lines[start_index : end_index + 1] = new_lines

        # 合并新内容并写入
        new_content = '\n'.join(lines)
        return self.writeDoc(user, doc_int, new_content)
    @rollback_on_exception
    def writeDoc(self, user: str, doc_int: int, content: str) -> Dict[str, Union[str,int]]:
        i_docId_Dict = self.tree.buildMdTree(user, doc_int, content, by="agent")
        if int(i_docId_Dict.get("success", -1)) < 0:
            return i_docId_Dict
        DocMap = self.session.query(affmap).where(affmap.intId == doc_int).first()# type: ignore
        if DocMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"DocMap from writeDoc,intId={doc_int} is None",doc_int)
        a_Dict = self.affine.updateDoc(user,DocMap.a_spaceId, DocMap.a_docId, content)
        a_Dict["nodes"]=self.tree.pNodeList([DocMap.intId],style="ja")
        return self.err.IF_SUCCESS_COMMIT_A2I(a_Dict, i_docId_Dict)
    @rollback_on_exception
    def moveDoc(self, user: str, doc_int: int, dest_space_int: int) -> Dict[str, Union[str,int]]:
        if doc_int==1:
            return self.err.INVALID_PARAM(f"moveDoc to root space is not allowed")
        i_docId_Dict = self.tree.moveNode(user, doc_int, dest_space_int, by="agent")
        if int(i_docId_Dict.get("success", -1)) < 0:
            return i_docId_Dict
        from_DocMap = self.session.query(affmap).where(affmap.intId == doc_int).first()# type: ignore
        dest_SpaceMap = self.session.query(affmap).where(affmap.intId == dest_space_int).first()# type: ignore
        if from_DocMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"from_DocMap from moveDoc,intId={doc_int} is None",doc_int)
        if dest_SpaceMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"dest_SpaceMap from moveDoc,intId={dest_space_int} is None",doc_int)
        dest_docMap=self.session.query(affmap).where(affmap.intId==doc_int).first()# type: ignore
        if dest_docMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"dest_docMap from moveDoc,intId={doc_int} is None",doc_int)
        a_dest_docId_Dict = self.affine.moveDoc(user,from_DocMap.a_spaceId, from_DocMap.a_docId, dest_SpaceMap.a_spaceId)
        dest_docMap.a_spaceId=dest_SpaceMap.a_spaceId
        dest_docMap.a_docId=a_dest_docId_Dict.get("affDocId", "")
        self.session.add(dest_docMap)
        a_dest_docId_Dict["nodes"]=self.tree.pNodeList([from_DocMap.intId,dest_SpaceMap.intId,dest_docMap.intId],style="ja")
        return self.err.IF_SUCCESS_COMMIT_A2I(a_dest_docId_Dict, i_docId_Dict)
    @rollback_on_exception

    def copyDoc(self, user: str, doc_int: int, dest_space_int: int) -> Dict[str, Union[str,int]]:
        if doc_int==1:
            return self.err.INVALID_PARAM(f"copyDoc to root space is not allowed")
        i_docId_Dict = self.tree.copyNode(user, doc_int, dest_space_int, by="agent")
        if int(i_docId_Dict.get("success", -1)) < 0:
            return i_docId_Dict
        from_DocMap = self.session.query(affmap).where(affmap.intId == doc_int).first()# type: ignore
        dest_SpaceMap = self.session.query(affmap).where(affmap.intId == dest_space_int).first()# type: ignore
        if from_DocMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"from_DocMap from copyDoc,intId={doc_int} is None",doc_int)
        if dest_SpaceMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"dest_SpaceMap from copyDoc,intId={dest_space_int} is None",doc_int)
        dest_docMap = affmap(intId=int(i_docId_Dict.get("success", -1)),a_spaceId=dest_SpaceMap.a_spaceId,a_docId="")
        self.session.add(dest_docMap)
        self.session.flush()
        a_dest_docId_Dict = self.affine.copyDoc(user,from_DocMap.a_spaceId, from_DocMap.a_docId, dest_SpaceMap.a_spaceId)
        dest_docMap.a_docId=a_dest_docId_Dict.get("affDocId", "")
        a_dest_docId_Dict["nodes"]=self.tree.pNodeList([from_DocMap.intId,dest_SpaceMap.intId,dest_docMap.intId],style="ja")
        return self.err.IF_SUCCESS_COMMIT_A2I(a_dest_docId_Dict, i_docId_Dict)

    def readDoc(self, user:str, doc_int:int,lineCnt:bool=False,by:str="agent")->Dict[str, Union[str,int]]:
        try:
            self.tree.permission_in_the_loop(user, doc_int, "read", by=by, processType=["doc"])
        except ValueError as e:
            return format_error(e)
        DocMap=self.session.query(affmap).where(affmap.intId==doc_int).first()# type: ignore
        if DocMap is None:
            return self.err.MAPFRAME_NOT_FOUND(f"DocMap from readDoc,intId={doc_int} is None",doc_int)
        doc_Dict = self.affine.readDoc(user,DocMap.a_spaceId, DocMap.a_docId)
        doc_Dict["success"]=doc_int
        if lineCnt:
            if doc_Dict["content"].strip()=="":
                doc_Dict["content"]=""
                return doc_Dict
            lines=doc_Dict["content"].split("\n")
            for i in range(len(lines)):
                lines[i]=str(i+1)+"|"+lines[i]
            doc_Dict["content"]="\n".join(lines)
        doc_Dict["nodes"]=self.tree.pNodeList([DocMap.intId],summaryContent="0:50",style="ja")
        return doc_Dict

    def search(self, user: str, query: str, ensure_full_return_type: str, method: str,subrootid:int=1) -> Dict[str, Union[str,int]]:
        treestr=self.tree.RAG(user,query,ensure_full_return_type,method,"agent","ta",subrootid=subrootid,surpressPrint=True)
        jsonstr=self.tree.RAG(user,query,ensure_full_return_type,method,"agent","ja",subrootid=subrootid,surpressPrint=True)
        return {"success":1,"content":treestr,"nodes":jsonstr,"message":"success"}
    @rollback_on_exception
    def constructTreeFromFs(self, user: str, rootPath: str):
        """
        从本地文件系统扫描指定的 rootPath，映射到affine后端的存储。测试用。根据以下约定进行简化：
        1. 根目录下只存在若干子目录（无文件，只有一层目录结构）。
        2. 每个子目录下仅包含文件（不再包含子目录）。
        
        步骤：
        1. 为 rootPath 创建一个空间，作为根目录对应的节点。
        2. 在根目录下遍历所有子目录，为每个子目录创建一个空间。
        3. 在每个子目录中遍历所有文件，为每个文件创建对应的文档。
        目前affine后端存储也近似用本地另一个文件夹代替。
        """
        for d in os.listdir(rootPath):
            subPath = os.path.join(rootPath, d)
            if os.path.isdir(subPath):
                print("found dir:",subPath)
                # 对每个子目录创建空间
                subDirId_Dict = self.createSpace(user, d)
                subDirId=subDirId_Dict.get("success",-1)
                if subDirId<0:
                    print(subDirId_Dict)
                    continue
                for f in os.listdir(subPath):
                    fileFullPath = os.path.join(subPath, f)
                    if os.path.isfile(fileFullPath):
                        print("found file:",fileFullPath)
                        docId_Dict = self.createDoc(user, subDirId, f)
                        docId=docId_Dict.get("success",-1)
                        if docId<0:
                            print(docId_Dict)
                            continue
                        with open(fileFullPath, "r", encoding="utf-8") as file:
                            content = file.read()
                        self.writeDoc(user, docId, content)
        return {"success": 1, "message": f"constructTreeFromFs 扫描完成，已处理路径：{rootPath}"}





